from strategy import strategy_base
from utils import date_util

# 策略1，最简单的贪心策略
# 在持有__max_keep_stock_num只股票时
# 每__max_keep_duration天调仓一次
# 并买入__recent_day天中涨幅最大的几只股票
class GreedyStrategy(strategy_base.TradeStrategyBase):
    def __init__(self, recent_day, max_keep_duration, max_keep_stock_num, init_money):
        strategy_base.TradeStrategyBase.__init__(self, max_keep_duration, max_keep_stock_num, init_money)
        self.recent_day = recent_day
        self.stock_array={}

    def buy_strategy(self, **kwargs):
        global unit_price, buy_count
        unit_price, buy_count = 0, 0
        time_data, now_date, start_date = kwargs["time_data"], kwargs["now_date"], kwargs["start_date"]
        # 股票太多了，不能买
        if len(self.stock_names) != self.max_keep_stock_num:
            stock_list = time_data.get_recent_max_increase(now_date, self.recent_day,
                                                       self.max_keep_stock_num - len(self.stock_names))
            if len(stock_list) == 0:
                return
            index = date_util.date_str_diff(now_date, start_date)
            unit_price = sum(time_data.stock_prices[stock_list[i]][index] for i in range(len(stock_list)))
            buy_count = int(self.money_now / unit_price)
            # 钱够，平均买入
            if buy_count != 0:
                for stock in stock_list:
                    self.stock_names.append([stock, 0, buy_count])
                    print("{0}: 买入了 {1} 股票 {2} 股".format(now_date, stock, buy_count))
                #self.stock_array.append(self.stock_names)
            # 钱不够，只买涨幅最大的
            else:
                unit_price = time_data.stock_prices[stock_list[0]][index]
                buy_count = int(self.money_now / unit_price)
                self.stock_names.append([stock_list[0], 0, buy_count])
                print("{0}: 买入了 {1} 股票 {2} 股".format(now_date, stock_list[0], buy_count))
        for i in range(0, len(self.stock_names)):
            self.stock_names[i][1] += 1
        self.stock_array[now_date]=self.stock_names
        self.money_now -= (unit_price * buy_count)
        print("{0}: 执行了买入操作，现在持有的股票状态：{1}, 剩余金钱： {2}".format(now_date, self.stock_names, self.money_now))

    def sell_strategy(self, **kwargs):
        time_data, now_date, start_date, end_flag = kwargs["time_data"], kwargs["now_date"], kwargs["start_date"], kwargs["end_flag"]
        # 到了调仓日期，卖股票
        index = date_util.date_str_diff(now_date, start_date)
        if end_flag:
            stocks_to_be_sold = self.stock_names
        else:
            stocks_to_be_sold = list(filter(lambda x: x[1] >= self.max_keep_duration, self.stock_names))
        if len(stocks_to_be_sold) == 0:
            return
        for stock in stocks_to_be_sold:
            print("{0}： 卖出了{1}, 每股 {2}".format(now_date, stock[0], time_data.stock_prices[stock[0]][index]))
            self.money_now += (stock[2] * time_data.stock_prices[stock[0]][index])
        if end_flag:
            self.stock_names = []
            #self.stock_array.append(self.stock_names)
        else:
            self.stock_names = list(filter(lambda x: x[1] < self.max_keep_duration, self.stock_names))
            self.stock_array[now_date]=self.stock_names
        print("{0}: 执行了卖出操作，现在持有的股票状态：{1}, 剩余金钱： {2}".format(now_date, self.stock_names, self.money_now))

